import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.datasets import load_linnerud
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score

# 1. 加载数据集
data = load_linnerud()
X = pd.DataFrame(data.data, columns=data.feature_names)  # 运动指标（特征）
y = pd.Series(data.target[:, 1], name="Weight")  # 体重作为目标变量

# 2. 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 3. 训练随机森林回归模型
rf_regressor = RandomForestRegressor(n_estimators=100, max_depth=5, random_state=42)
rf_regressor.fit(X_train, y_train)

# 4. 进行预测
y_pred = rf_regressor.predict(X_test)

# 5. 计算性能指标
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)

print(f"均方误差 (MSE): {mse:.2f}")
print(f"R² 分数: {r2:.2f}")

# 6. 可视化：实际值 vs 预测值
plt.figure(figsize=(8, 6))
plt.scatter(y_test, y_pred, color="blue", alpha=0.6, label="预测值")
plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', label="理想情况")
plt.xlabel("实际体重")
plt.ylabel("预测体重")
plt.title("随机森林回归：实际 vs 预测")
plt.legend()
plt.show()
